# backend/model_client.py
"""
模型服务客户端
封装与模型推理服务的 HTTP 通信
"""
import os
import requests
from typing import Dict, Any, Optional
import time

# 模型服务配置
MODEL_SERVICE_URL = os.environ.get("MODEL_SERVICE_URL", "http://10.143.12.79:8252")
MODEL_API_TOKEN = os.environ.get("MODEL_API_TOKEN", None)
REQUEST_TIMEOUT = int(os.environ.get("MODEL_REQUEST_TIMEOUT", "300"))  # 5分钟超时

def check_model_service_health() -> bool:
    """检查模型服务是否健康"""
    try:
        response = requests.get(
            f"{MODEL_SERVICE_URL}/health",
            timeout=5
        )
        if response.status_code == 200:
            data = response.json()
            return data.get("model_loaded", False) and data.get("tokenizer_loaded", False)
        return False
    except Exception as e:
        print(f"⚠️ [模型服务] 健康检查失败: {e}")
        return False

def analyze_contract(
    source_code: str,
    vuln_type: str = "Reentrancy",
    max_new_tokens: int = 256,
    temperature: float = 0.7,
    top_p: float = 0.9,
    do_sample: bool = False
) -> Dict[str, Any]:
    """
    调用模型服务进行合约分析
    
    Args:
        source_code: 合约源代码
        vuln_type: 漏洞类型
        max_new_tokens: 最大生成token数
        temperature: 温度参数
        top_p: top_p参数
        do_sample: 是否采样
    
    Returns:
        分析结果字典，格式与 single_detect.analyze_contract 相同
    """
    url = f"{MODEL_SERVICE_URL}/api/analyze"
    
    payload = {
        "source_code": source_code,
        "vuln_type": vuln_type,
        "max_new_tokens": max_new_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "do_sample": do_sample
    }
    
    headers = {
        "Content-Type": "application/json"
    }
    
    # 如果配置了 API Token，添加到请求头
    if MODEL_API_TOKEN:
        headers["Authorization"] = f"Bearer {MODEL_API_TOKEN}"
    
    try:
        print(f"📡 [模型服务] 发送分析请求到 {url}")
        start_time = time.time()
        
        response = requests.post(
            url,
            json=payload,
            headers=headers,
            timeout=REQUEST_TIMEOUT
        )
        
        elapsed_time = time.time() - start_time
        
        if response.status_code == 200:
            result = response.json()
            print(f"✅ [模型服务] 分析完成，耗时 {elapsed_time:.2f}秒")
            
            # 转换为与 single_detect.analyze_contract 相同的格式
            return {
                'detected': result.get('detected', False),
                'status': result.get('status', 'Unknown'),
                'analysis': result.get('analysis', ''),
                'has_vulnerability': result.get('has_vulnerability', False)
            }
        elif response.status_code == 503:
            error_msg = "模型服务未就绪，请稍后重试"
            print(f"❌ [模型服务] {error_msg}")
            raise RuntimeError(error_msg)
        else:
            error_msg = f"模型服务返回错误: HTTP {response.status_code}"
            try:
                error_detail = response.json().get('detail', '')
                error_msg += f" - {error_detail}"
            except:
                error_msg += f" - {response.text[:200]}"
            print(f"❌ [模型服务] {error_msg}")
            raise RuntimeError(error_msg)
            
    except requests.exceptions.Timeout:
        error_msg = f"模型服务请求超时（>{REQUEST_TIMEOUT}秒）"
        print(f"❌ [模型服务] {error_msg}")
        raise RuntimeError(error_msg)
    except requests.exceptions.ConnectionError:
        error_msg = f"无法连接到模型服务: {MODEL_SERVICE_URL}"
        print(f"❌ [模型服务] {error_msg}")
        raise RuntimeError(error_msg)
    except Exception as e:
        error_msg = f"模型服务调用失败: {str(e)}"
        print(f"❌ [模型服务] {error_msg}")
        raise RuntimeError(error_msg)


